"""DHF模型工具模块

提供数据处理、结果分析、可视化等辅助功能。
"""

from typing import List, Dict, Any

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from loguru import logger

from util.font_config import setup_chinese_font

import warnings
import matplotlib

# 彻底抑制matplotlib的所有警告
warnings.filterwarnings('ignore')
matplotlib.set_loglevel('CRITICAL')  # 只显示严重错误

# 抑制特定的字体警告
warnings.filterwarnings('ignore', message='.*Glyph.*missing from font.*')
warnings.filterwarnings('ignore', message='.*findfont.*')
warnings.filterwarnings('ignore', category=UserWarning, module='matplotlib.*')

# 设置matplotlib后端为非交互式（减少警告）
matplotlib.use('Agg', force=True)

class ModelUtils:
    """DHF模型工具类"""

    @staticmethod
    def prepare_input_data(rainfall_data: List[float],
                           evaporation_data: List[float] = None,
                           time_step: float = 1.0) -> pd.DataFrame:
        """
        准备模型输入数据
        
        Args:
            rainfall_data: 降雨数据序列 (mm)
            evaporation_data: 蒸发数据序列 (mm)，如果为None则使用默认值
            time_step: 时间步长 (小时)
            
        Returns:
            包含输入数据的DataFrame
        """
        if evaporation_data is None:
            evaporation_data = [3.0] * len(rainfall_data)  # 默认蒸发量

        if len(rainfall_data) != len(evaporation_data):
            raise ValueError("降雨数据和蒸发数据长度不一致")

        # 创建时间索引
        time_index = pd.date_range(
            start='2023-01-01',
            periods=len(rainfall_data),
            freq=f'{int(time_step)}H'
        )

        df = pd.DataFrame({
            'rainfall': rainfall_data,
            'evaporation': evaporation_data,
            'time_step': time_step
        }, index=time_index)

        return df

    @staticmethod
    def calculate_model_performance(observed: np.ndarray,
                                    simulated: np.ndarray) -> Dict[str, float]:
        """
        计算模型性能指标
        
        Args:
            observed: 观测值数组
            simulated: 模拟值数组
            
        Returns:
            性能指标字典
        """
        if len(observed) != len(simulated):
            min_len = min(len(observed), len(simulated))
            observed = observed[:min_len]
            simulated = simulated[:min_len]

        # 去除无效值
        valid_mask = ~(np.isnan(observed) | np.isnan(simulated))
        obs = observed[valid_mask]
        sim = simulated[valid_mask]

        if len(obs) == 0:
            return {'NSE': 0, 'RMSE': 0, 'MAE': 0, 'R2': 0, 'PBIAS': 0}

        # Nash-Sutcliffe效率系数
        nse = 1 - np.sum((obs - sim) ** 2) / np.sum((obs - np.mean(obs)) ** 2)

        # 均方根误差
        rmse = np.sqrt(np.mean((obs - sim) ** 2))

        # 平均绝对误差
        mae = np.mean(np.abs(obs - sim))

        # 决定系数
        corr_coef = np.corrcoef(obs, sim)[0, 1]
        r2 = corr_coef ** 2 if not np.isnan(corr_coef) else 0

        # 百分比偏差
        pbias = 100 * np.sum(sim - obs) / np.sum(obs) if np.sum(obs) != 0 else 0

        return {
            'NSE': nse,  # Nash-Sutcliffe效率系数
            'RMSE': rmse,  # 均方根误差
            'MAE': mae,  # 平均绝对误差
            'R2': r2,  # 决定系数
            'PBIAS': pbias  # 百分比偏差
        }

    # 在文件顶部添加导入
    import sys
    import os
    sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

    # 在所有绘图方法中，删除原来的字体设置，替换为：
    # 例如在 plot_hydrograph 方法中：
    @staticmethod
    def plot_hydrograph(time_series: pd.DatetimeIndex,
                        observed: np.ndarray = None,
                        simulated: np.ndarray = None,
                        rainfall: np.ndarray = None,
                        title: str = "流量过程线",
                        save_path: str = None) -> plt.Figure:
        """
        绘制流量过程线
        """
        # 确保字体配置正确
        setup_chinese_font()

        fig, ax1 = plt.subplots(figsize=(12, 8))

        # 绘制流量
        if observed is not None:
            ax1.plot(time_series[:len(observed)], observed, 'b-',
                     label='观测流量', linewidth=2)

        if simulated is not None:
            ax1.plot(time_series[:len(simulated)], simulated, 'r--',
                     label='模拟流量', linewidth=2)

        ax1.set_xlabel('时间')
        ax1.set_ylabel('流量 (m³/s)', color='b')
        ax1.tick_params(axis='y', labelcolor='b')
        ax1.grid(True, alpha=0.3)

        # 绘制降雨
        if rainfall is not None:
            ax2 = ax1.twinx()
            ax2.bar(time_series[:len(rainfall)], rainfall,
                    alpha=0.3, color='gray', label='降雨量')
            ax2.set_ylabel('降雨量 (mm)', color='gray')
            ax2.tick_params(axis='y', labelcolor='gray')
            ax2.invert_yaxis()  # 反转y轴，降雨从上往下

        # 添加图例
        lines1, labels1 = ax1.get_legend_handles_labels()
        if rainfall is not None:
            lines2, labels2 = ax2.get_legend_handles_labels()
            ax1.legend(lines1 + lines2, labels1 + labels2, loc='upper right')
        else:
            ax1.legend()

        plt.title(title, fontsize=14, fontweight='bold')
        plt.tight_layout()

        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')

        return fig

    @staticmethod
    def plot_runoff_components(time_series: pd.DatetimeIndex,
                               components: Dict[str, np.ndarray],
                               rainfall_data: np.ndarray = None,
                               title: str = "径流分量过程线",
                               save_path: str = None) -> plt.Figure:
        """
        绘制径流分量过程线
        
        Args:
            time_series: 时间序列
            components: 径流分量字典
            rainfall_data: 降雨数据数组 (mm)
            title: 图表标题
            save_path: 保存路径
            
        Returns:
            matplotlib图形对象
        """
        # 确保字体配置正确
        setup_chinese_font()
    
        fig, ax1 = plt.subplots(figsize=(12, 8))
    
        # 如果有降雨数据，创建第二个y轴用于显示降雨
        if rainfall_data is not None:
            ax2 = ax1.twinx()
            
            # 绘制降雨柱状图（倒置，从上往下）
            rainfall_aligned = rainfall_data[:len(time_series)]
            ax2.bar(time_series[:len(rainfall_aligned)], rainfall_aligned, 
                    color='lightblue', alpha=0.6, width=0.03, label='降雨')
            ax2.set_ylabel('降雨量 (mm)', color='blue')
            ax2.tick_params(axis='y', labelcolor='blue')
            ax2.invert_yaxis()  # 倒置y轴，降雨从上往下显示
            ax2.legend(loc='upper right')
    
        # 绘制径流分量线图
        colors = ['blue', 'green', 'orange', 'red']
        labels = {
            'impervious_runoff': '不透水面积径流',
            'interflow': '地面壤中流',
            'groundwater_runoff': '地下径流',
            'total_runoff': '总径流'
        }
    
        for i, (component, values) in enumerate(components.items()):
            if component in labels:
                ax1.plot(time_series[:len(values)], values,
                        color=colors[i % len(colors)],
                        label=labels[component],
                        linewidth=2)
    
        ax1.set_xlabel('时间')
        ax1.set_ylabel('径流深 (mm)')
        ax1.grid(True, alpha=0.3)
        ax1.legend(loc='upper left')
    
        plt.title(title, fontsize=14, fontweight='bold')
        plt.tight_layout()
    
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
    
        return fig

    @staticmethod
    def plot_storage_dynamics(time_series: pd.DatetimeIndex,
                              surface_storage: np.ndarray,
                              subsurface_storage: np.ndarray,
                              title: str = "蓄水量动态变化",
                              save_path: str = None) -> plt.Figure:
        """
        绘制蓄水量动态变化图
        
        Args:
            time_series: 时间序列
            surface_storage: 表层蓄水量 (mm)
            subsurface_storage: 下层蓄水量 (mm)
            title: 图表标题
            save_path: 保存路径
            
        Returns:
            matplotlib图形对象
        """
        # 确保字体配置正确
        setup_chinese_font()

        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 10), sharex=True)

        # 表层蓄水量
        ax1.plot(time_series[:len(surface_storage)], surface_storage,
                 'b-', linewidth=2, label='表层蓄水量')
        ax1.set_ylabel('表层蓄水量 (mm)')
        ax1.grid(True, alpha=0.3)
        ax1.legend()

        # 下层蓄水量
        ax2.plot(time_series[:len(subsurface_storage)], subsurface_storage,
                 'g-', linewidth=2, label='下层蓄水量')
        ax2.set_xlabel('时间')
        ax2.set_ylabel('下层蓄水量 (mm)')
        ax2.grid(True, alpha=0.3)
        ax2.legend()

        plt.suptitle(title, fontsize=14, fontweight='bold')
        plt.tight_layout()

        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')

        return fig

    @staticmethod
    def plot_scatter_comparison(observed: np.ndarray,
                                simulated: np.ndarray,
                                title: str = "观测值与模拟值对比",
                                save_path: str = None) -> plt.Figure:
        """
        绘制观测值与模拟值散点对比图
        
        Args:
            observed: 观测值数组
            simulated: 模拟值数组
            title: 图表标题
            save_path: 保存路径
            
        Returns:
            matplotlib图形对象
        """
        # 确保字体配置正确
        setup_chinese_font()

        # 确保数组长度一致
        min_len = min(len(observed), len(simulated))
        obs = observed[:min_len]
        sim = simulated[:min_len]

        # 去除无效值
        valid_mask = ~(np.isnan(obs) | np.isnan(sim))
        obs = obs[valid_mask]
        sim = sim[valid_mask]

        fig, ax = plt.subplots(figsize=(8, 8))

        # 绘制散点图
        ax.scatter(obs, sim, alpha=0.6, s=30)

        # 绘制1:1线
        min_val = min(np.min(obs), np.min(sim))
        max_val = max(np.max(obs), np.max(sim))
        ax.plot([min_val, max_val], [min_val, max_val], 'r--', linewidth=2, label='1:1线')

        # 计算性能指标
        performance = ModelUtils.calculate_model_performance(obs, sim)

        # 添加性能指标文本
        textstr = f'NSE = {performance["NSE"]:.3f}\n'
        textstr += f'R² = {performance["R2"]:.3f}\n'
        textstr += f'RMSE = {performance["RMSE"]:.2f}\n'
        textstr += f'PBIAS = {performance["PBIAS"]:.1f}%'

        props = dict(boxstyle='round', facecolor='wheat', alpha=0.8)
        ax.text(0.05, 0.95, textstr, transform=ax.transAxes, fontsize=10,
                verticalalignment='top', bbox=props)

        ax.set_xlabel('观测值')
        ax.set_ylabel('模拟值')
        ax.grid(True, alpha=0.3)
        ax.legend()

        plt.title(title, fontsize=14, fontweight='bold')
        plt.tight_layout()

        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')

        return fig

    @staticmethod
    def export_results_to_excel(results: Dict[str, Any],
                                file_path: str,
                                time_series: pd.DatetimeIndex = None):
        """
        将结果导出到Excel文件
        
        Args:
            results: 结果字典
            file_path: 文件保存路径
            time_series: 时间序列
        """
        try:
            with pd.ExcelWriter(file_path, engine='openpyxl') as writer:

                # 导出时间序列数据
                if time_series is not None and 'flow_series' in results:
                    flow_series = results['flow_series']
                    min_length = min(len(time_series), len(flow_series))

                    df_flow = pd.DataFrame({
                        '时间': time_series[:min_length],
                        '流量(m³/s)': flow_series[:min_length]
                    })
                    df_flow.to_excel(writer, sheet_name='流量过程', index=False)

                # 导出径流分量
                if 'runoff_components' in results:
                    components = results['runoff_components']
                    if components:  # 确保components不为空
                        max_len = max(len(v) for v in components.values() if len(v) > 0)

                        df_components = pd.DataFrame()
                        if time_series is not None:
                            df_components['时间'] = time_series[:max_len]

                        component_names = {
                            'impervious_runoff': '不透水面积径流(mm)',
                            'interflow': '地面壤中流(mm)',
                            'groundwater_runoff': '地下径流(mm)',
                            'total_runoff': '总径流(mm)'
                        }

                        for key, values in components.items():
                            if key in component_names and len(values) > 0:
                                extended_values = np.zeros(max_len)
                                extended_values[:len(values)] = values
                                df_components[component_names[key]] = extended_values

                        if not df_components.empty:
                            df_components.to_excel(writer, sheet_name='径流分量', index=False)

                # 导出性能指标
                if 'performance' in results and results['performance']:
                    df_performance = pd.DataFrame([
                        {'指标': '纳什效率系数(NSE)', '值': results['performance']['NSE']},
                        {'指标': '决定系数(R²)', '值': results['performance']['R2']},
                        {'指标': '均方根误差(RMSE)', '值': results['performance']['RMSE']},
                        {'指标': '平均绝对误差(MAE)', '值': results['performance']['MAE']},
                        {'指标': '百分比偏差(PBIAS)', '值': results['performance']['PBIAS']}
                    ])
                    df_performance.to_excel(writer, sheet_name='性能指标', index=False)

                # 确保至少有一个工作表
                if len(writer.sheets) == 0:
                    # 创建一个基本的摘要工作表
                    df_summary = pd.DataFrame({
                        '项目': ['模型类型', '计算状态'],
                        '值': ['DHF模型', '已完成']
                    })
                    df_summary.to_excel(writer, sheet_name='模型摘要', index=False)

        except Exception as e:
            logger.error(f"Excel导出失败: {e}")
            raise

    @staticmethod
    def create_model_summary_report(results: Dict[str, Any]) -> str:
        """
        创建模型运行总结报告
        
        Args:
            results: 模型结果字典
            
        Returns:
            总结报告字符串
        """
        report = "DHF模型运行总结报告\n"
        report += "=" * 50 + "\n\n"

        # 基本信息
        if 'peak_info' in results:
            peak_info = results['peak_info']
            report += f"洪峰流量: {peak_info['peak_flow']:.2f} m³/s\n"
            report += f"峰现时间: {peak_info['peak_time']:.1f} 小时\n\n"

        # 径流总量
        if 'total_volume' in results:
            report += f"径流总量: {results['total_volume']:.2f} 万m³\n\n"

        # 性能指标
        if 'performance' in results:
            perf = results['performance']
            report += "模型性能指标:\n"
            report += f"  纳什效率系数(NSE): {perf['NSE']:.3f}\n"
            report += f"  决定系数(R²): {perf['R2']:.3f}\n"
            report += f"  均方根误差(RMSE): {perf['RMSE']:.2f}\n"
            report += f"  平均绝对误差(MAE): {perf['MAE']:.2f}\n"
            report += f"  百分比偏差(PBIAS): {perf['PBIAS']:.1f}%\n\n"

        # 模型评价
        if 'performance' in results:
            nse = results['performance']['NSE']
            if nse > 0.75:
                evaluation = "优秀"
            elif nse > 0.65:
                evaluation = "良好"
            elif nse > 0.50:
                evaluation = "满意"
            else:
                evaluation = "不满意"

            report += f"模型整体评价: {evaluation}\n"

        return report
